import os
from stable_baselines3 import PPO
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.callbacks import CheckpointCallback
from env_turtlebot3 import TurtleBot3Env

def main():
    # 日志目录
    log_dir = "./ppo_tb3_log/"
    os.makedirs(log_dir, exist_ok=True)

    # 包装环境
    env = TurtleBot3Env()
    env = Monitor(env, log_dir)

    # 继续加载旧模型
    model = PPO.load("ppo_turtlebot3", env=env, tensorboard_log=log_dir)

    # checkpoint 保存设置
    checkpoint_callback = CheckpointCallback(
        save_freq=10000,
        save_path="./ppo_tb3_checkpoints/",
        name_prefix="ppo_tb3"
    )

    # 继续训练
    model.learn(total_timesteps=50000, callback=checkpoint_callback)

    # 覆盖保存
    model.save("ppo_turtlebot3")

if __name__ == "__main__":
    main()

